import numpy as np
import pykgraph


def ivecs_read(fname):
    a = np.fromfile(fname, dtype='int32')
    d = a[0]
    return a.reshape(-1, d + 1)[:, 1:].copy(), d


def fvecs_read(fname):
    data, d = ivecs_read(fname)
    return data.view('float32').astype(np.float32), d


def save_graph(data, ds_name):
    index = pykgraph.KGraph(data, 'euclidean')
    index.build(reverse=-1)
    path = "../../graph/KGraph/%s.graph" % ds_name
    index.save_text(path)


if __name__ == "__main__":
    dataset_name = 'siftsmall'
    base_path = "../../data/%s/base.fvecs" % dataset_name
    base, d = fvecs_read(base_path)
    save_graph(base, dataset_name)
